#include "vtkMeshExport.hpp"
#include "vtkConstants.hpp"
#include "vtkFunctions.hpp"
#include "vtkFunctions_export.hpp"
#include "meshImportExport.hpp"
#include "meshTypes.hpp"
#include "streamAccess.hpp"

#include <iomanip>
#include <tuple>

namespace hms::vtk {

namespace xml = tinyxml2;
namespace fs = std::filesystem;

template<typename MeshType>
std::string wholeExtent( const StructMeshBase<MeshType>& mesh ){
	std::stringstream converter { getConverter() };
	
	converter
		<<  "0 " << mesh.nCellsX()
		<< " 0 " << mesh.nCellsY()
		<< " 0 0";
	return converter.str();
}

xml::XMLElement* coordsToDataArray(
	xml::XMLElement* parent,
	const Ref<const VectorXs>& coords
){
	xml::XMLElement* dataArray { createDataArray(parent) };
	if ( coords.size() > 1 ){
		dataArray->SetAttribute( "RangeMin", "0" );
		dataArray->SetAttribute(
			"RangeMax", std::to_string(coords.size()-1).c_str()
		);
	}

	std::stringstream converter { getConverter() };
	write(converter, coords);

	dataArray->SetText( converter.str().c_str() );

	return dataArray;
}

xml::XMLElement* verticesToDataArray(
	xml::XMLElement* parent,
	const Ref<const Matrix2Xs>& vertices
){
	xml::XMLElement* dataArray { createDataArray(parent) };
	dataArray->SetAttribute( nComponentsName, "3" );

	std::stringstream converter { getConverter() };
	for ( const Ref<const Matrix2Xs>::ConstColXpr col : vertices.colwise() )
		write(converter, col) << " 0 ";

	dataArray->SetText( converter.str().c_str() );

	return dataArray;
}

xml::XMLElement* createPoints(
	xml::XMLElement* parent,
	const Ref<const Matrix2Xs>& vertices
){
	xml::XMLElement* points = parent->InsertNewChildElement("Points");

	verticesToDataArray( points, vertices );

	return points;
}

template<typename MeshType>
xml::XMLElement* pieceWithWholeExtent(
	xml::XMLElement* data,
	const StructMeshBase<MeshType>& mesh
){
	data->SetAttribute( "WholeExtent", wholeExtent(mesh).c_str() );

	xml::XMLElement* piece = data->InsertNewChildElement("Piece");
	piece->SetAttribute( "Extent", wholeExtent(mesh).c_str() );
	return piece;
}

void meshToVtk( xml::XMLElement* data, const UniMesh& mesh ){
	std::stringstream converter { getConverter() };
	
	pieceWithWholeExtent(data,mesh);

	/* origin */
	converter.str({});
	write( converter, mesh.origin() ) << " 0";
	data->SetAttribute( "Origin", converter.str().c_str() );

	/* cell spacing */
	converter.str({});
	converter << mesh.dx() << " " << mesh.dy() << " 0";
	data->SetAttribute( "Spacing", converter.str().c_str() );
}

void meshToVtk( xml::XMLElement* data, const RectMesh& mesh ){
	std::stringstream converter { getConverter() };

	xml::XMLElement* piece { pieceWithWholeExtent(data,mesh) };

	xml::XMLElement* coords = piece->InsertNewChildElement("Coordinates");

	/* only write the inner coordinates, not the ones for ghost cells */
	coordsToDataArray( coords, mesh.xCoordsInner() );
	coordsToDataArray( coords, mesh.yCoordsInner() );
	coordsToDataArray( coords, VectorNs<1>{0} );
}

void meshToVtk( xml::XMLElement* data, const StructMesh& mesh ){
	std::stringstream converter { getConverter() };

	xml::XMLElement* piece { pieceWithWholeExtent(data,mesh) };

	createPoints( piece, mesh.vertices() );
}

/**
 * @brief Writes Mesh data to a VTK XML document
 * 
 * @param data The XML document's first child node
 * @param mesh The hms::Mesh object to be written
 */
void meshToVtk(
	xml::XMLElement* data,
	const Mesh& mesh
){
	/* create data node depending on requested representation type */
	using namespace std::string_literals;
	xml::XMLElement* piece { data->InsertNewChildElement("Piece") };

	piece->SetAttribute(
		"NumberOfPoints", std::to_string( mesh.nVertices() ).c_str()
	);
	piece->SetAttribute(
		( usePolyData ? "NumberOfPolys"s : "NumberOfCells"s ).c_str(),
		std::to_string( mesh.nDomainCells() ).c_str()
	);
	if (usePolyData){
		piece->SetAttribute("NumberOfVerts" ,"0");
		piece->SetAttribute("NumberOfLines" ,"0");
		piece->SetAttribute("NumberOfStrips","0");
	}

	createPoints( piece, mesh.vertices() );

	xml::XMLElement* cells { piece->InsertNewChildElement(
		( usePolyData ? "Polys"s : "Cells"s ).c_str()
	) };
	xml::XMLElement* connectivity { createDataArray( cells, "Int64") };
	xml::XMLElement* offsets      { createDataArray( cells, "Int64") };

	connectivity->SetAttribute( "Name", "connectivity" );
	offsets->SetAttribute( "Name", "offsets" );

	{
		std::stringstream converter { getConverter() };
		for (Index i = 0; i<mesh.nDomainCells(); ++i){
			write( converter, mesh.cell(i).vertices(true) );
			converter << " ";
		}
		
		connectivity->SetText( converter.str().c_str() );
	}{
		ArrayXi offset (mesh.nDomainCells() );
		offset(0) = mesh.cell(0).nEdges();
		for (Index i = 1; i<offset.size(); ++i)
			offset(i) = offset(i-1) + mesh.cell(i).nEdges();
		
		std::stringstream converter { getConverter() };
		write(converter, offset);

		offsets->SetText( converter.str().c_str() );
	}
	
	if (!usePolyData){
		xml::XMLElement* types { createDataArray( cells, "UInt8") };
		types->SetAttribute( "Name", "types" );
		std::stringstream converter { getConverter() };
		/* the VTK format defines linear cell types, with type 5 being a
		 * triangle and type 9 being a quadrilateral */
		for (Index i = 0; i<mesh.nDomainCells(); ++i)
			converter << (mesh.cell(i).nEdges() == 3 ? 5 : 9) << " ";
		
		types->SetText( converter.str().c_str() );
	}
}

template<typename MeshType>
void meshToVtk( xml::XMLDocument* doc, const MeshBase<MeshType>& mesh ){
	/* create the document and the elements which are equal for all mesh types*/
	constexpr auto typeStr { getVtkFormat<MeshType>() };

	xml::XMLElement* vtk { createVtkNode(doc) };
	vtk->SetAttribute( "type", typeStr );
	xml::XMLElement* dataElem { vtk->InsertNewChildElement(typeStr) };
	
	meshToVtk( dataElem, mesh.derived() );
}


void meshToVtk(
	xml::XMLDocument* doc,
	const MeshVariant& meshVar
){
	std::visit( [&](const auto& mesh){
		meshToVtk(doc, mesh);
	}, meshVar );
}

std::unique_ptr<xml::XMLDocument> meshToVtk(
	const MeshVariant& meshVar
){
	auto doc { std::make_unique<xml::XMLDocument>() };

	meshToVtk( doc.get(), meshVar );

	return doc;
}


void meshToVtk(
	const fs::path& path, const hms::MeshVariant& meshVar, bool boundary
){
	std::unique_ptr<xml::XMLDocument> doc { meshToVtk(meshVar) };
	toFile(doc.get(), path, meshVar);

	if (boundary)
		boundaryToXML( path.parent_path(), meshVar );
}

void meshToVtk(
	const fs::path& caseDir,
	const std::string& caseName,
	const MeshVariant& meshVar,
	bool boundary
){
	fs::path meshPath {caseDir / caseName};
	replaceExtension(meshPath, meshVar);
	meshToVtk(meshPath, meshVar, boundary);
}


void boundaryToXML( xml::XMLNode* node, const MeshVariant& meshVar ){
	std::visit( [&](const auto& mesh){

		xml::XMLElement* boundary {
			getRootOrInsert(node, boundaryNode)
		};

		for ( const BoundaryPatch&
			patch : mesh.boundaryPatches()
		){
			xml::XMLElement* bp { findChildWithAttributeOrInsert(
				boundary, boundaryChild, "Name", patch.name().c_str()
			)};

			/* continuous boundary patches can be represented by three
			 * vertices: first, second (for direction) and last */
			Array3i verts { boundarySegmentToVertices(
				mesh,
				patch.edges()
			) };
			xml::XMLElement* vertsEl { findOrInsert( bp, "Vertices" ) };
			matrixToText( vertsEl,verts);
		}
	}, meshVar);
}

std::unique_ptr<xml::XMLDocument> boundaryToXML( const MeshVariant& meshVar ){
	auto doc { std::make_unique<xml::XMLDocument>() };

	boundaryToXML( doc.get(), meshVar );
	return doc;
}

void boundaryToXML( const fs::path& caseDir, const MeshVariant& meshVar ){
	toFile( boundaryToXML(meshVar).get(), getBoundaryPath(caseDir) );
}

void removeNonexistentPatches( xml::XMLNode* node, const MeshVariant& meshVar ){
	xml::XMLElement* boundary { getBoundaryNode(node) };

	std::visit( [&](const auto& mesh){
		xml::XMLElement* patchElem { boundary->FirstChildElement("Patch") };
		while (patchElem){
			bool meshHasPatch {false};
			for ( const BoundaryPatch& patch : mesh.boundaryPatches() ){
				if ( patchElem->Attribute( "Name", patch.name().c_str() ) ){
					meshHasPatch = true;
					break;
				}
			}
			if ( !meshHasPatch ){
				std::cout << "Removing patch \"" << patchElem->Attribute("Name") << "\"...\n";
				xml::XMLElement* toDelete {patchElem};
				patchElem = patchElem->NextSiblingElement("Patch");
				boundary->DeleteChild(toDelete);
				continue;
			}
			patchElem = patchElem->NextSiblingElement("Patch");
		}
	}, meshVar );
}

#define INSTANTIATE_FOR_ALL_MESHTYPES(MeshType) \
template void meshToVtk( xml::XMLDocument*, const MeshBase<MeshType>& );
#include "meshInstantiator.hpp"

} // end namespace hms